package cn.jackyu.stock.stock_grid_trading.event.producer;

import org.springframework.http.MediaType;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.Map;
import java.util.function.Consumer;

public class SseClientUtil {

    private final WebClient webClient;
    private final Map<String, Disposable> subscriptionMap; // 存储股票代码和对应的Disposable

    public SseClientUtil(WebClient.Builder webClientBuilder, String baseUrl) {
        this.webClient = webClientBuilder.build();
        this.subscriptionMap = new HashMap<>();
    }

    /**
     * 订阅SSE流
     *
     * @param code           股票代码（secid）
     * @param messageHandler 消息处理函数
     */
    public void subscribeFirst(String code, Consumer<String> messageHandler) {
        String uri = "https://61.push2.eastmoney.com/api/qt/stock/trends2/sse?" +
                "fields1=f1,f2,f3,f4,f5,f6,f7,f8,f9,f10,f11,f12,f13,f17" +
                "&fields2=f51,f52,f53,f54,f55,f56,f57,f58" +
                "&mpi=1000&ut=fa5fd1943c7b386f172d6893dbfba10b" +
                "&secid=" + code +
                "&ndays=1&iscr=0&iscca=0&wbp2u=|0|0|0|web";

        // 订阅SSE流
        Flux<String> eventStream = webClient.get()
                .uri(uri)
                .header("Referer", "https://quote.eastmoney.com/gb/zsKLSE.html") // 设置Referer
                .header("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/134.0.0.0 Safari/537.36 Edg/134.0.0.0") // 设置User-Agent
                .accept(MediaType.TEXT_EVENT_STREAM)
                .retrieve()
                .bodyToFlux(String.class);

        // 存储Disposable对象
        Disposable disposable = eventStream.subscribe(
                messageHandler::accept,
                error -> System.err.println("Error occurred for code " + code + ": " + error.getMessage()),
                () -> System.out.println("SSE stream completed for code: " + code)
        );

        // 将股票代码和Disposable存入Map
        subscriptionMap.put(code, disposable);
    }

    /**
     * 订阅SSE流
     *
     * @param code           股票代码（secid）
     * @param messageHandler 消息处理函数
     */
    public void subscribe(String code, Consumer<String> messageHandler) {
        String uri = "https://5.push2.eastmoney.com/api/qt/stock/details/sse?fields1=f1,f2,f3,f4&fields2=f51,f52,f53,f54,f55&mpi=1000&dect=1" +
                "&ut=bd1d9ddb04089700cf9c27f6f7426281&fltt=2&pos=-11" +
                "&secid=" + code +
                "&wbp2u=|0|0|0|web";

        // 订阅SSE流
        Flux<String> eventStream = webClient.get()
                .uri(uri)
                .header("Referer", "https://quote.eastmoney.com/gb/zsKLSE.html") // 设置Referer
                .header("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/134.0.0.0 Safari/537.36 Edg/134.0.0.0") // 设置User-Agent
                .accept(MediaType.TEXT_EVENT_STREAM)
                .retrieve()
                .bodyToFlux(String.class);

        // 存储Disposable对象
        Disposable disposable = eventStream.subscribe(
                messageHandler::accept,
                error -> System.err.println("Error occurred for code " + code + ": " + error.getMessage()),
                () -> System.out.println("SSE stream completed for code: " + code)
        );

        // 将股票代码和Disposable存入Map
        subscriptionMap.put(code, disposable);
    }

    /**
     * 根据股票代码取消订阅
     *
     * @param code 股票代码
     */
    public void unsubscribe(String code) {
        Disposable disposable = subscriptionMap.get(code);
        if (disposable != null && !disposable.isDisposed()) {
            disposable.dispose(); // 取消订阅
            subscriptionMap.remove(code); // 从Map中移除
            System.out.println("Unsubscribed for code: " + code);
        } else {
            System.out.println("No active subscription found for code: " + code);
        }
    }

    /**
     * 获取所有订阅的股票代码
     *
     * @return 股票代码列表
     */
    public Iterable<String> getSubscribedCodes() {
        return subscriptionMap.keySet();
    }

    /**
     * 关闭所有订阅
     */
    public void unsubscribeAll() {
        for (Map.Entry<String, Disposable> entry : subscriptionMap.entrySet()) {
            String code = entry.getKey();
            Disposable disposable = entry.getValue();
            if (!disposable.isDisposed()) {
                disposable.dispose();
                System.out.println("Unsubscribed for code: " + code);
            }
        }
        subscriptionMap.clear(); // 清空Map
        System.out.println("All subscriptions have been closed.");
    }
}